
#include <cusp/precond/detail/device/spai.h>
#include <cusp/multiply.h>

namespace cusp
{
namespace precond
{

/* Create with same sparsity pattern as A */
template <typename ValueType, typename MemorySpace>
template <typename MatrixType>
spai<ValueType,MemorySpace>::spai(const MatrixType& A)
	: linear_operator<ValueType,MemorySpace>(A.num_rows, A.num_cols, A.num_entries)
{
	make_spai(A, A, typename MatrixType::memory_space(),
			typename MatrixType::format(),
			typename MatrixType::format());
}


/* Create with sparsity pattern P */
template <typename ValueType, typename MemorySpace>
template <typename MatrixType, typename MatrixType2>
spai<ValueType,MemorySpace>::spai(const MatrixType& A, const MatrixType2& P)
	: linear_operator<ValueType,MemorySpace>(P.num_rows, P.num_cols, P.num_entries)
{
	make_spai(A, P, typename MatrixType::memory_space(),
			typename MatrixType::format(),
			typename MatrixType2::format());
}


/** A is stored in host memory. Move it to device, and converting if needed. */
template <typename ValueType, typename MemorySpace>
template <typename MatrixType, typename MatrixType2>
void spai<ValueType,MemorySpace>::make_spai(const MatrixType& A, const MatrixType2& P, cusp::host_memory, cusp::known_format, cusp::known_format)
{
	matrix_type Ad = A;
	make_spai(Ad, P, typename matrix_type::memory_space(),
			 typename matrix_type::format(),
			 typename MatrixType2::format());
}


/** A is in device memory in the right format. P is in the right format. */
template <typename ValueType, typename MemorySpace>
template <typename MatrixType, typename MatrixType2>
void spai<ValueType,MemorySpace>::make_spai(const MatrixType& A, const MatrixType2& P, cusp::device_memory, cusp::csr_format, cusp::csr_format)
{
	/* resize the result matrix and set the sparsity pattern */
	matrix.resize(P.num_rows, P.num_cols, P.num_entries);
	thrust::copy(P.row_offsets.begin(), P.row_offsets.end(), matrix.row_offsets.begin());
	thrust::copy(P.column_indices.begin(), P.column_indices.end(), matrix.column_indices.begin());

	cusp::precond::detail::device::spai_csr<MatrixType>(A, matrix);
}


/** A and P are on device but (at least one) need to be converted. */
template <typename ValueType, typename MemorySpace>
template <typename MatrixType, typename MatrixType2>
void spai<ValueType,MemorySpace>::make_spai(const MatrixType& A, const MatrixType2& P, cusp::device_memory, cusp::known_format, cusp::known_format)
{
	matrix = P;
	matrix_type Acsr = A;
	cusp::precond::detail::device::spai_csr<matrix_type>(Acsr, matrix);
}


template <typename ValueType, typename MemorySpace>
template <typename VectorType1, typename VectorType2>
void spai<ValueType,MemorySpace>::operator()(const VectorType1& x, VectorType2& y) const
{
	cusp::multiply(matrix, x, y);
}


} // end namespace precond
} // end namespace cusp

